/**
 * Copyright (C) 2010 - 2013 Harry Glasgow
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package com.googlecode.jaden.ui.dialog.manualtraining;

import com.googlecode.jaden.common.config.ProjectConfiguration;
import com.googlecode.jaden.common.config.LayerConfig;
import com.googlecode.jaden.common.other.JadenException;
import com.googlecode.jaden.common.other.Utils;
import com.googlecode.jaden.common.enums.LayerType;
import com.googlecode.jaden.engine.*;
import com.googlecode.jaden.ui.UiUtils;
import static com.googlecode.jaden.ui.Constants.*;
import com.googlecode.jaden.ui.UI;
import com.googlecode.jaden.ui.dialog.BaseDialog;
import org.jfree.chart.ChartFactory;
import org.jfree.chart.ChartPanel;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.plot.PlotOrientation;
import org.jfree.data.xy.XYSeries;
import org.jfree.data.xy.XYSeriesCollection;

import javax.swing.*;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.util.Map;
import java.util.HashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.prefs.Preferences;

public class ManualTrainingDialog extends BaseDialog {

    private final UI ui;
    private final ProjectConfiguration project;
    private final Preferences preferences;

    private final Map<Integer, Layer> layers = new HashMap<Integer, Layer>();
    private double manualLearningRate;
    private double manualMomentum;
    private int manualEpochs;

    private final AtomicBoolean running = new AtomicBoolean();
    private final JButton configureButton;
    private final JButton startStopButton;
    private final JButton resetButton;
    private final JButton jogButton;
    private final JButton analyzeButton;

    private final XYSeries trainingTimeSeries;
    private final XYSeries validationTimeSeries;

    private final ExecutorService service;

    public ManualTrainingDialog(UI ui, Preferences preferences, ProjectConfiguration project) {
        super(ui, DIALOG_MANUAL_TRAINING_TITLE);
        this.ui = ui;
        this.project = project;
        this.preferences = preferences;
        if (project.getBestLayers().isEmpty()) {
            createLayers();
        }
        layers.putAll(project.getBestLayers());
        manualLearningRate = project.getManualLearningRate();
        manualMomentum = project.getManualMomentum();
        manualEpochs = project.getManualEpochs();
        trainingTimeSeries = new XYSeries(Utils.translate(LABEL_TRAINING_TEXT));
        validationTimeSeries = new XYSeries(Utils.translate(LABEL_VALIDATION_TEXT));
        XYSeriesCollection xySeriesCollection = new XYSeriesCollection();
        xySeriesCollection.addSeries(trainingTimeSeries);
        xySeriesCollection.addSeries(validationTimeSeries);
        JFreeChart jfc = ChartFactory.createXYLineChart(project.getName(),  Utils.translate(LABEL_EPOCHS_TEXT),
                Utils.translate(LABEL_ERROR_TEXT), xySeriesCollection, PlotOrientation.VERTICAL, true, true, true);
        ChartPanel cp = new ChartPanel(jfc);
        JPanel mainPanel = UiUtils.verticalList(cp);
        setMainPanel(mainPanel, preferences);

        configureButton = UiUtils.createButton(LABEL_CONFIGURE_TEXT);
        addButton(configureButton);
        configureButton.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e) {
                doManualConfiguration();
            }
        });

        startStopButton = UiUtils.createButton(LABEL_START_TEXT);
        addButton(startStopButton);
        startStopButton.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e) {
                doStartStop();
            }
        });

        resetButton = UiUtils.createButton(LABEL_RESET_TEXT);
        addButton(resetButton);
        resetButton.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e) {
                for (Layer layer : layers.values()) {
                    layer.reset();
                }
                trainingTimeSeries.clear();
                validationTimeSeries.clear();
            }
        });

        jogButton = UiUtils.createButton(LABEL_JOG_TEXT);
        addButton(jogButton);
        jogButton.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e) {
                for (Layer layer : layers.values()) {
                    layer.jog();
                }
            }
        });

        analyzeButton = UiUtils.createButton(LABEL_ANALYZE_TEXT);
        addButton(analyzeButton);
        analyzeButton.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e) {
                doAnalysis();
            }
        });

        JButton dataButton = UiUtils.createButton(LABEL_DATA_TEXT);
        addButton(dataButton);
        dataButton.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e) {
                doData();
            }
        });

        service = Executors.newFixedThreadPool(1);
    }

    private void doData() {
        try {
            DataSet set = new DataSet(project);
            Network network = new Network(layers, manualLearningRate, manualMomentum);
            Model model = new Model(project.getInputConfigurations(), network, project.getOutputConfiguration());

            String[][] trainingInputValues = set.getTrainingInputValues();
            String[] trainingOutputValues = set.getTrainingOutputValues();
            String[] trainingActualValues = new String[trainingOutputValues.length];

            String[][] validationInputValues = set.getValidationInputValues();
            String[] validationOutputValues = set.getValidationOutputValues();
            String[] validationActualValues = new String[validationOutputValues.length];

            String[][] testingInputValues = set.getTestingInputValues();
            String[] testingOutputValues = set.getTestingOutputValues();
            String[] testingActualValues = new String[testingOutputValues.length];

            for (int i = 0; i < trainingOutputValues.length; i++) {
                String[] input = trainingInputValues[i];
                String actual = model.forward(input);
                trainingActualValues[i] = actual;
            }
            for (int i = 0; i < validationOutputValues.length; i++) {
                String[] input = validationInputValues[i];
                String actual = model.forward(input);
                validationActualValues[i] = actual;
            }
            for (int i = 0; i < testingOutputValues.length; i++) {
                String[] input = testingInputValues[i];
                String actual = model.forward(input);
                testingActualValues[i] = actual;
            }

            TestResultsDialog dialog = new TestResultsDialog(this, preferences, project.getInputConfigurations(),
                    project.getOutputConfiguration(),
                    trainingInputValues, trainingOutputValues, trainingActualValues,
                    validationInputValues, validationOutputValues, validationActualValues,
                    testingInputValues, testingOutputValues, testingActualValues
            );
            dialog.setVisible(true);
        } catch (Exception e) {
            Utils.handleException(e);
            JOptionPane.showMessageDialog(this, Utils.translate(ERROR_GENERAL));
        }
    }

    // Create a basic two-layer (I + O) / 2 : O network.
    private void createLayers() {
        int inputWidth = project.getInternalInputWidth();
        int outputWidth = project.getInternalOutputWidth();
        int middleWidth = (inputWidth + outputWidth) / 2;
        LayerConfig config1 = new LayerConfig(LayerType.Normal, middleWidth);
        LayerConfig config2 = new LayerConfig(LayerType.Normal, outputWidth);
        Layer layer1 = new Layer(config1, inputWidth);
        Layer layer2 = new Layer(config2, middleWidth);
        layers.put(0, layer1);
        layers.put(1, layer2);
    }

    private void doAnalysis() {
        AnalysisDialog dialog = new AnalysisDialog(ui, getPreferences(), project.getInputConfigurations(), 
                layers, project.getOutputConfiguration());
        dialog.setVisible(true);
    }

    private void doManualConfiguration() {
        new ConfigureManualTrainingDialog(this, getPreferences(), layers, manualLearningRate, manualMomentum,
                manualEpochs, project.getInternalInputWidth(), project.getInternalOutputWidth()).setVisible(true);
    }

    protected void doOk() {
        ui.updateManualTraining(manualLearningRate, manualMomentum, manualEpochs, layers);
        dispose();
    }

    private void doStartStop() {
        synchronized (running) {
            running.set(!running.get());
            if (running.get()) {
                okButton.setEnabled(false);
                cancelButton.setEnabled(false);
                configureButton.setEnabled(false);
                startStopButton.setText(Utils.translate(LABEL_STOP_TEXT));
                resetButton.setEnabled(false);
                jogButton.setEnabled(false);
                analyzeButton.setEnabled(false);
                service.execute(new NetworkRunner());
            } else {
                okButton.setEnabled(true);
                cancelButton.setEnabled(true);
                configureButton.setEnabled(true);
                startStopButton.setText(Utils.translate(LABEL_START_TEXT));
                resetButton.setEnabled(true);
                jogButton.setEnabled(true);
                analyzeButton.setEnabled(true);
            }
        }
    }

    public void updateManualTraining(Map<Integer, Layer> layersArg, double learningRate, double momentum, int epochs) {
        layers.clear();
        layers.putAll(layersArg);
        manualLearningRate = learningRate;
        manualMomentum = momentum;
        manualEpochs = epochs;
    }

    ///////////////////
    // Inner Classes //
    ///////////////////

    private class NetworkRunner implements Runnable {
        public void run() {
            try {
                DataSet ds = new DataSet(project);
                trainingTimeSeries.clear();
                validationTimeSeries.clear();
                Network network = new Network(layers, manualLearningRate, manualMomentum);
                Model model = new Model(project.getInputConfigurations(), network, project.getOutputConfiguration());
                Trainer trainer = new Trainer(ds, model);
                long epoch = 0;
                while (running.get()) {
                    double trainingError = trainer.train();
                    double validationError = trainer.validate();
                    trainingTimeSeries.add(epoch, trainingError);
                    validationTimeSeries.add(epoch, validationError);
                    epoch++;
                    if (epoch > manualEpochs) {
                        if (running.get()) {
                            doStartStop();
                        }
                    }
                }
            } catch (JadenException e) {
                doStartStop();
                Utils.handleException(e);
                JOptionPane.showMessageDialog(ManualTrainingDialog.this, e.getMessage());
            } catch (Exception e) {
                doStartStop();
                Utils.handleException(e);
                JOptionPane.showMessageDialog(ManualTrainingDialog.this, Utils.translate(ERROR_GENERAL));
            }
        }
    }
}